"""
    用于对pth权重文件进行测试
    输入为训练时所用的同尺寸的小图
    add by likang
"""
# -*- coding: utf-8 -*-
"""
道路裂缝检测与修复系统 - 基于深度学习的裂缝分割模型
核心功能：使用预训练模型对道路图像进行裂缝检测，输出裂缝概率图

输入为训练时所用的同尺寸的小图
"""

import os
import cv2  # OpenCV库，用于图像处理
import numpy as np  # 数值计算库
import torch  # PyTorch深度学习框架
from matplotlib import pyplot as plt
from model.deepcrack_with_uu_Sp_3 import DeepCrack_uu_Sp_3  # 自定义的裂缝分割模型架构
from tqdm import tqdm  # 进度条显示工具
from trainer.trainer import DeepCrackTrainer  # 模型训练器
from config.config import Config as cfg  # 配置文件（包含模型参数、训练参数等）

# 设置CUDA可见设备（指定使用的GPU）
# os.environ["CUDA_VISIBLE_DEVICES"] = cfg.gpu_id


def repair(input_path, output_path,save_in_local=True):
    """
    主处理函数 - 执行裂缝检测全流程

    参数:
        input_path: str - 输入图像目录路径
        output_path: str - 结果保存目录路径
    """

    # ==================== 1. 模型初始化配置 ====================
    # 预训练模型权重路径（注意：路径需根据实际位置修改）
    pretrained_model = r'predicted_model/DuSp3_CRKWH1001_Wight6_2_200_DuSp3_CRKWH1001_Wight6_2_200_epoch(199)_0000200_2025-05-08-13-25-51.pth'

    # 创建输出目录（如果不存在）
    if not os.path.exists(output_path):
        os.makedirs(output_path)  # 递归创建目录

    # ==================== 2. 模型加载与配置 ====================
    # 设置计算设备（优先使用GPU）
    device = torch.device("cuda")

    # 初始化模型架构
    model = DeepCrack_uu_Sp_3()  # 使用配置文件参数构建模型

    # 多GPU并行计算配置（device_ids指定使用的GPU索引）
    model = torch.nn.DataParallel(model, device_ids=[0]).cuda()

    # 初始化训练器并加载预训练权重
    trainer = DeepCrackTrainer(model,False).cuda()
    model.load_state_dict(
        # torch.load(pretrained_model, map_location='gpu'),
        trainer.saver.load(pretrained_model, multi_gpu=True),
        strict=False  # 严格匹配权重名称
    )
    model.eval()  # 设置为评估模式（关闭dropout等训练专用层）

    # ==================== 3. 图像处理流水线 ====================
    with torch.no_grad():  # 禁用梯度计算（节省内存）
        # 获取输入目录下所有文件
        files = os.listdir(input_path)
        # 按照文件名排序
        files.sort()
        # 没有文件就退出
        if not files:
            print("输入目录为空，请检查路径！")
            return

        # 获取前pre_num张图像进行检测（默认检测前10张）
        pre_num = 1
        # 使用进度条遍历处理最后pre_num张图像
        # for file in tqdm(files[-pre_num:], desc="裂缝检测进度"):# 这是从后往前

        for file in tqdm(files[:pre_num], desc="裂缝检测进度"): # 这是从前往后
            print("正在处理：", file)
            # 3.1 文件路径处理
            save_path = os.path.normpath(os.path.join(output_path, file))  # 规范化路径
            img_path = os.path.join(input_path, file)  # 输入图像完整路径

            # 3.2 图像读取与预处理
            image = cv2.imread(img_path)  # 读取BGR格式图像
            w, h, c = image.shape  # 获取图像尺寸（宽度、高度、通道数）
            # print(image.shape) #(256, 256, 3)
            image_tensor = image.transpose(2, 0, 1)  # HWC -> CHW（PyTorch标准格式）
            # print(image_tensor.shape) #(3 ,256, 256)
            # print(image_tensor)  #0-255范围
            image_tensor = image_tensor.astype(np.float32) / 255.0  # 归一化到[0,1]
            # print(image_tensor)  # 0-1范围

            # 转换为PyTorch张量并添加batch维度     0-1范围+batch维度 共同作为模型的输入
            test_data = torch.from_numpy(image_tensor).float().unsqueeze(0).to(device)
            """
            test_pred是6 tensor的元组 取第一个全尺度进行展示
            """
            test_pred = model(test_data)  # 前向传播
            test_pred = torch.sigmoid(test_pred[0].cpu().squeeze())  # 压缩维度(即将1维度去掉)并应用sigmoid
            # print(test_pred)
            # print(test_pred.shape) #torch.Size([256, 256])

            #显示预测图
            # plt.figure(figsize=(w / 100, h / 100))
            # plt.imshow(test_pred)
            # plt.title('test_pred_'+str(file))
            # plt.show()

            # ==================== 4. 结果后处理与保存 ====================
            """
            将预测结果保存为PNG图像。
            对预测结果进行后处理（取0.7次方以增强对比度）。
            将概率图转换为0-255的灰度图并保存。
            """
            # 生成输出文件名（保持原文件名，扩展名改为.png）
            save_name = os.path.join(output_path, file.split('.')[0] + '.png')

            # 概率图增强（gamma校正提升对比度）
            pree = np.power(test_pred.data.numpy(), 0.7)  # 指数0.7为经验值

            # 转换为0-255范围的灰度图
            save_result = (test_pred.data.numpy() * 255).astype(np.uint8)
            # plt显示 原图尺寸
            plt.figure(figsize=(w/100, h/100))
            plt.title('test_pred_'+str(file))
            plt.imshow(save_result, cmap='gray')
            plt.show()

            # 用户自己选择是否保存结果图像
            if save_in_local:
                # 保存结果图像
                cv2.imwrite(save_name, save_result)




if __name__ == '__main__':
    """
    输入: 一组道路图像，训练时用的小尺寸图，如256*256。
    输出: 对应的裂缝检测结果（灰度图，白色表示裂缝）。
    模型: 基于DeepCrack和DANet的改进模型，用于裂缝检测。
    """
    # 输入输出路径配置（示例路径，实际使用时需修改）
    input_path = r'/home/likang/DeepCrack_LK/evaluate/data/validation/images'  # 输入图像目录
    output_path = r'./test_data/out_small'  # 结果输出目录

    # 执行裂缝检测流程
    repair(input_path, output_path,save_in_local=False)